package com.ideaaedi.shardingdynamictable.config.sharding.support;

import com.ideaaedi.shardingdynamictable.entity.bo.DatasourceArgs;
import com.ideaaedi.shardingdynamictable.entity.bo.UpdateActualDataNodeBO;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.algorithm.config.AlgorithmProvidedShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import java.util.TreeSet;
import java.util.stream.Collectors;

/**
 * 刷新actual-data-nodes支持类 <br /> <br /> 参考：https://github.com/apache/shardingsphere/issues/16725 <br /> <br />
 * 说明：我们在配置文件中指定了数据存储节点涉及到了很多表，那是为以后作考虑的；现在项目中真实存在的表可能没那么多, 所以我们需要一个自动根据数据库中已存在的分表，设置对应的actual-data-nodes数据节点的功能. <br />
 * <br /> 注：如果不设置的话，当进行查询（如：全查）时，就会因为表不存在而报错
 *
 * @author JustryDeng
 * @since 2022/10/7 16:06:22
 */
@Slf4j
@Component
public class RefreshActualDataNodesSupport {
    
    @Resource
    ShardingSphereDataSource shardingSphereDataSource;
    
    /**
     * 刷新actual-data-nodes
     */
    public void refresh(UpdateActualDataNodeBO bo) {
        final String logicTableName = bo.getLogicTableName();
        if (StringUtils.isBlank(logicTableName)) {
            throw new IllegalArgumentException("logicTableName cannot be null.");
        }
        final List<Pair<DatasourceArgs, TreeSet<String>>> dsAndRealTablesPairList = bo.getDsAndRealTablesPairList();
        if (CollectionUtils.isEmpty(dsAndRealTablesPairList)) {
            throw new IllegalArgumentException("dsAndRealTablesPairList cannot be empty.");
        }
        StringBuilder actualDataNodesBuilder = new StringBuilder();
        for (Pair<DatasourceArgs, TreeSet<String>> datasourceArgsTreeSetPair : dsAndRealTablesPairList) {
            final DatasourceArgs datasourceArgs = datasourceArgsTreeSetPair.getLeft();
            Objects.requireNonNull(datasourceArgs, "DatasourceArgs cannot be null.");
            final String datasourceName = datasourceArgs.getName();
            if (StringUtils.isBlank(datasourceName)) {
                throw new IllegalArgumentException("datasourceName cannot be null.");
            }
            final TreeSet<String> realTableNames = datasourceArgsTreeSetPair.getRight();
            if (StringUtils.isBlank(datasourceName)) {
                throw new IllegalArgumentException("datasourceName cannot be null.");
            }
            if (CollectionUtils.isEmpty(realTableNames)) {
                throw new IllegalArgumentException("realTableNames cannot be empty.");
            }
            final String realTableNameInline =
                    realTableNames.stream().map(realTableName -> "'" + realTableName + "'").collect(Collectors.joining(","
                    ));
            // String actualDataNodes = "ds-0.t_order_$->{[1..8]}"; // 等价于ds-0.t_order_${[1,2,3,4,5,6,7,8]}
            // 多数据源的话，使用逗号分割, 如:: "ds-0.t_order_$->{[1..8]},ds-1.t_order_123}"
            actualDataNodesBuilder.append(datasourceName).append(".${[").append(realTableNameInline).append("]}").append(",");
        }
        // 去除尾巴上的逗号
        actualDataNodesBuilder.replace(actualDataNodesBuilder.length() - 1, actualDataNodesBuilder.length(), "");
        String actualDataNodes = actualDataNodesBuilder.toString();
        log.info("refresh actual-data-nodes to " + actualDataNodes);
        updateShardRuleActualDataNodes(shardingSphereDataSource, shardingSphereDataSource.getSchemaName(), logicTableName, actualDataNodes);
    }
    
    /**
     * 刷新actual-data-nodes
     */
    private void updateShardRuleActualDataNodes(ShardingSphereDataSource dataSource, String schemaName,
                                                String logicTableName, String newActualDataNodes) {
        // Context manager.
        ContextManager contextManager = dataSource.getContextManager();
        
        // Rule configuration.
        Collection<RuleConfiguration> newRuleConfigList = new LinkedList<>();
        Collection<RuleConfiguration> oldRuleConfigList = contextManager
                .getMetaDataContexts()
                .getMetaData(schemaName)
                .getRuleMetaData()
                .getConfigurations();
        
        for (RuleConfiguration oldRuleConfig : oldRuleConfigList) {
            if (oldRuleConfig instanceof AlgorithmProvidedShardingRuleConfiguration) {
                
                // Algorithm provided sharding rule configuration
                AlgorithmProvidedShardingRuleConfiguration oldAlgorithmConfig =
                        (AlgorithmProvidedShardingRuleConfiguration) oldRuleConfig;
                AlgorithmProvidedShardingRuleConfiguration newAlgorithmConfig =
                        new AlgorithmProvidedShardingRuleConfiguration();
                
                // Sharding table rule configuration Collection
                Collection<ShardingTableRuleConfiguration> newTableRuleConfigList = new LinkedList<>();
                Collection<ShardingTableRuleConfiguration> oldTableRuleConfigList = oldAlgorithmConfig.getTables();
                
                oldTableRuleConfigList.forEach(oldTableRuleConfig -> {
                    if (logicTableName.equals(oldTableRuleConfig.getLogicTable())) {
                        // 在旧的nodes上进行补充（旧的nodes与新的nodes可能有部分重复，这里不管，因为sharding会自动去重， 见org.apache.shardingsphere.sharding.rule.TableRule#addActualTable）
                        String mergeActualDataNodes;
                        ///if (StringUtils.isNotBlank(oldTableRuleConfig.getActualDataNodes())) {
                        ///    mergeActualDataNodes = oldTableRuleConfig.getActualDataNodes() + "," + newActualDataNodes;
                        ///} else {
                        mergeActualDataNodes = newActualDataNodes;
                        ///}
                        ShardingTableRuleConfiguration newTableRuleConfig =
                                new ShardingTableRuleConfiguration(oldTableRuleConfig.getLogicTable(), mergeActualDataNodes);
                        newTableRuleConfig.setTableShardingStrategy(oldTableRuleConfig.getTableShardingStrategy());
                        newTableRuleConfig.setDatabaseShardingStrategy(oldTableRuleConfig.getDatabaseShardingStrategy());
                        newTableRuleConfig.setKeyGenerateStrategy(oldTableRuleConfig.getKeyGenerateStrategy());
                        
                        newTableRuleConfigList.add(newTableRuleConfig);
                    } else {
                        newTableRuleConfigList.add(oldTableRuleConfig);
                    }
                });
                
                newAlgorithmConfig.setTables(newTableRuleConfigList);
                newAlgorithmConfig.setAutoTables(oldAlgorithmConfig.getAutoTables());
                newAlgorithmConfig.setBindingTableGroups(oldAlgorithmConfig.getBindingTableGroups());
                newAlgorithmConfig.setBroadcastTables(oldAlgorithmConfig.getBroadcastTables());
                newAlgorithmConfig.setDefaultDatabaseShardingStrategy(oldAlgorithmConfig.getDefaultDatabaseShardingStrategy());
                newAlgorithmConfig.setDefaultTableShardingStrategy(oldAlgorithmConfig.getDefaultTableShardingStrategy());
                newAlgorithmConfig.setDefaultKeyGenerateStrategy(oldAlgorithmConfig.getDefaultKeyGenerateStrategy());
                newAlgorithmConfig.setDefaultShardingColumn(oldAlgorithmConfig.getDefaultShardingColumn());
                newAlgorithmConfig.setShardingAlgorithms(oldAlgorithmConfig.getShardingAlgorithms());
                newAlgorithmConfig.setKeyGenerators(oldAlgorithmConfig.getKeyGenerators());
                
                newRuleConfigList.add(newAlgorithmConfig);
            }
        }
        
        // update context
        contextManager.alterRuleConfiguration(schemaName, newRuleConfigList);
    }
}
